# encoding=utf-8

"""JavaWebSocket.py
    建立与Ruo-yi管理端的链接

    增加等待队列，来管理需要发送的信息
"""
import json
import threading
import time
import traceback

import websocket
from websocket import WebSocketApp
from concurrent.futures.thread import ThreadPoolExecutor

from seg_common.TaskQueue.ITaskQueue import ITaskQueue
from seg_common.Socket.IWebSocketClient import IWebSocketClient
from seg_common.LoggingService import ConsoleService


class JavaUploadMsgQueue(ITaskQueue):
    def __init__(self, worker_size: int = 4, max_buffer_size: int = 1000, **kwargs):
        super(JavaUploadMsgQueue, self).__init__(**kwargs)

        # 容器
        self.max_buffer_size = max_buffer_size
        self.task_queue_handler = self.make_storage()

        # 写入线程池, 保证主程序可以正常运行
        self.write_executor = ThreadPoolExecutor(max_workers=worker_size)

        # 互斥锁
        self.lock = threading.Lock()

    def make_storage(self, **kwargs):
        return []

    def task_in_async(self, o, **kwargs):
        future = self.write_executor.submit(self.task_in, o, **kwargs)
        future.add_done_callback(self.thread_callback)

    def thread_callback(self, worker):
        worker_exception = worker.exception()
        if worker_exception:
            msg = 'Task in thread error occur: {}'.format(worker_exception)
            ConsoleService.console_log(msg, ConsoleService.EXCEPTION)
            ConsoleService.console_log("Please check code, it should not occur", ConsoleService.EXCEPTION)

    def task_in(self, o, **kwargs):
        # TODO 最起码要保证系统的稳定运行，这里舍弃一些数据的精度
        # TODO 后续改成写入文本的形式
        self.lock.acquire()
        if len(self.task_queue_handler) > self.max_buffer_size:
            self.task_queue_handler.pop(0)
        self.lock.release()

        self.task_queue_handler.append(o)

    def task_out(self, **kwargs):
        if len(self.task_queue_handler) == 0:
            return None
        return self.task_queue_handler.pop(0)


class JavaWebSocket(IWebSocketClient):
    def __init__(self, uuid: str = None, **kwargs):
        super(JavaWebSocket, self).__init__(uuid, **kwargs)

        # 等待队列
        self.msg_queue = JavaUploadMsgQueue()
        # 上传执行线程
        self.upload_executor = ThreadPoolExecutor(max_workers=2)
        # 启动循环
        self.upload_executor.submit(self.upload_loop)

    def upload_to_task_queue(self, msg: dict, **kwargs):
        json_str = json.dumps(msg)
        self.msg_queue.task_in_async(json_str)

    def upload_to_java(self, msg: str, **kwargs):
        while True:
            if self.get_connect_status():
                self.ws.send(msg)
                break
            else:
                time.sleep(0.1)

    def thread_callback(self, worker):
        worker_exception = worker.exception()
        if worker_exception:
            msg = 'upload thread error occur: {}'.format(worker_exception)
            ConsoleService.console_log(msg, ConsoleService.EXCEPTION)
            ConsoleService.console_log("Please check code, it should not occur", ConsoleService.EXCEPTION)

    def push(self, msg: dict, **kwargs):
        try:
            self.upload_to_task_queue(msg, **kwargs)
        except Exception as e:
            print(traceback.format_exc())
            return -1, 'json convert fail'
        return 0, 'upload success'

    def upload_loop(self):
        while True:
            task = self.msg_queue.task_out()

            if task is None:
                time.sleep(0.2)
            else:
                future = self.upload_executor.submit(self.upload_to_java, task)
                future.add_done_callback(self.thread_callback)

    def msg_arrive_default(self, ws: WebSocketApp, msg):
        self.log.info('Java msg arrive: {}'.format(msg))


if __name__ == '__main__':
    client = JavaWebSocket()
    thread2 = threading.Thread(target=client.start)
    thread2.start()

    d = {}
    i = 0
    # d['count'] = i
    # client.push(d)
    while True:
        d['count'] = i
        i += 1
        client.push(d)

        time.sleep(1)
